from pybrain.rl.experiments.episodic import EpisodicExperiment
from publicgood_agent import *
from publicgood_task import *
import numpy

class DelayedRewardExperiment(EpisodicExperiment):
    """A modification of EpisodicExperiment so that the agent has to wait until the others are done making
        decisions before they are rewarded"""
    def __init__(self, task, agent):
        EpisodicExperiment.__init__(self, task, agent)

    def _oneInteraction(self, verbosity):
        """ Give the observation to the agent, takes its resulting action and returns
            it to the task. Then gives the reward to the agent again and returns it.
        """
        if self.doOptimization:
            raise Exception('When using a black-box learning algorithm, only full episodes can be done.')
        else:
            self.agent.newEpisode()
            self.stepid += 1
            obs = self.task.observation
            if verbosity==2:
                print(str(self.agent)+" sees, "+str(obs))
            self.agent.integrateObservation(obs)
            act = self.task.performAction(self.agent.getAction())
            if verbosity==2:
                print(str(self.agent)+" does, "+str(act))

    def giveReward(self):
        reward = self.task.getReward()
        self.agent.giveReward(reward)
        return reward
    

class TwoTaskExperiment(EpisodicExperiment):
    """An experiment with tasks alternated per interaction"""
    def __init__(self, task, agent):
        EpisodicExperiment.__init__(self, task, agent)
        self.allocation = DelayedRewardExperiment(task.contribute, agent.allocator)
        if agent.punisher != None:
            self.punishment = DelayedRewardExperiment(task.punish, agent.punisher)
            self.punishing = True
        else:
            self.punishing = False
        self.t = 0

    def _oneInteraction(self, verbose):
        if self.t % 2 == 0 or not self.punishing:
            self.allocation._oneInteraction(verbose)
        else:
            self.punishment._oneInteraction(verbose)

        self.t += 1

    def giveReward(self,verbosity):
        allo = self.allocation.giveReward()
        if self.punishing:
            puni = self.punishment.giveReward()
            assert allo == puni
        if verbosity==2:
            print(str(self.agent)+" is rewarded, "+str(allo))
        
        
class PublicGoodExperiment(EpisodicExperiment):
    """Like many of the other PublicGood classes, this holds sub-experiment classes that are run/accessed through this container"""
    def __init__(self, PGE, allocator_type=1, punisher_type=1):
        """Since the number of players is defined in the PublicGoodEnvironment=PGE, this also optionally inits the agents"""
        print("Making "+str(PGE.nplayers)+" agents.")
        self.PGE = PGE
        self.stepid = 0

        self.agents = []
        PublicGoodAgent.index = 0
        for i in range(PGE.nplayers):
            self.agents.append(PublicGoodAgent())


        self.tasks = []
        for agent in self.agents:
            task = PublicGoodTask(PGE, agent)
            agent.setSubAgents(task, allocator_type, punisher_type)
            self.tasks.append(task)


        self.experiments = []
        for i in range(PGE.nplayers):
            self.experiments.append(TwoTaskExperiment(self.tasks[i], self.agents[i]))

    def _oneInteraction(self, verbosity=False):
        """Interactions for each sub-exp, learning, and output all in one method!"""        
        print("Round: " + str(self.stepid))

        #Allocation
        if verbosity==2:
            print("Allocation...")
        for task in self.tasks:
            task.contribute.getObservation()
        for exp in self.experiments:
            exp._oneInteraction(verbosity)
        if verbosity==1:
            print("Allocations:")
            print(self.PGE.allocations)
        #Punishment
        if self.PGE.epsilon != 0:
            if verbosity==2:
                print("Punishment...")
            for task in self.tasks:
                task.punish.getObservation()
            for exp in self.experiments:
                exp._oneInteraction(verbosity)
            if verbosity==1:
                print("Punishments:")
                print(self.PGE.punishments)
        #Learning
        if verbosity==2:
            print("Learning...")
        for exp in self.experiments:
            exp.giveReward(verbosity)
            
        for agent in self.agents:
            agent.learn()
##            agent.reset()

        #self.PGE.reset()
        self.stepid += 1


    def doEpisodes(self, number = 1, verbosity=1, measures=[]):

        args = [self.PGE.getAllocations]
        if self.PGE.epsilon != 0:
            args.insert(1,self.PGE.getPunishments)
            
        returns = numpy.zeros((len(measures), self.PGE.nplayers, number))
        
        for t in range(number):
            self._oneInteraction(verbosity)
            run_args = []
            for arg in args:
                run_args.append(arg())
            for m, measure in enumerate(measures):
                for n in range(self.PGE.nplayers):
                    returns[m,n,t] = measure(n, *run_args) 

        return returns
        
                    
